use crate::common::*;

use std::collections::{HashMap, HashSet};
use std::str;

use ordered_float::NotNan;
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyBytes, PyString};
use serde::Deserialize;

#[derive(Deserialize)]
struct BitmexRecordData {
    table: String,
    action: String,
    data: Vec<Entry>,
}

#[derive(Deserialize)]
struct Entry {
    symbol: String,
    id: i64,
    side: String,
    #[serde(default)]
    size: i64,
    #[serde(default)]
    price: NotNan<f64>,
}

#[pyclass]
pub struct BitmexHandler {
    books: HashMap<String, Py<BookBuilder>>,
    book_callbacks: Vec<PyObject>,
    book_symbols: HashSet<String>,
    id_to_price: HashMap<i64, NotNan<f64>>,
    any_snapshot: bool,
}

#[pymethods]
impl BitmexHandler {
    #[new]
    fn new_python(obj: &PyRawObject, _py: Python) {
        obj.init(BitmexHandler {
            books: HashMap::new(),
            book_symbols: HashSet::new(),
            book_callbacks: Vec::new(),
            id_to_price: HashMap::new(),
            any_snapshot: false,
        });
    }

    fn on_orderbook_l2(
        &mut self,
        py: Python,
        timestamp: i64,
        record_bytes: &PyBytes,
        _queue_data: &PyAny,
        _topic_data: &PyAny,
    ) -> PyResult<()> {
        let bytes: &[u8] = record_bytes.as_bytes();

        let record = Record::<BitmexRecordData> {
            timestamp: timestamp,
            data: serde_json::from_slice(bytes).expect("cannot parse json"),
        };

        // {'table': 'orderBookL2', 'action': 'update', 'data': [{'symbol': 'XBTUSD', 'id': 8799039200, 'side': 'Sell', 'size': 185942}]}
        assert!(record.data.table == "orderBookL2");
        let mut touched_symbols: HashSet<String> = HashSet::new();
        for entry in &record.data.data {
            if self.book_symbols.contains(&entry.symbol) {
                touched_symbols.insert(entry.symbol.clone());
            }
        }
        match &*record.data.action {
            "partial" => {
                self.any_snapshot = true;
                // {
                // "table":"orderBookL2_25",
                // "action":"partial",
                // "data":[
                //     {"symbol":"XBTUSD","id":17999992000,"side":"Sell","size":100,"price":80},
                //     {"symbol":"XBTUSD","id":17999993000,"side":"Sell","size":20,"price":70},
                //     {"symbol":"XBTUSD","id":17999994000,"side":"Sell","size":10,"price":60},
                //     {"symbol":"XBTUSD","id":17999995000,"side":"Buy","size":10,"price":50},
                //     {"symbol":"XBTUSD","id":17999996000,"side":"Buy","size":20,"price":40},
                //     {"symbol":"XBTUSD","id":17999997000,"side":"Buy","size":100,"price":30}
                // ]
                // }

                // clear all books.
                // all symbols come in one partial event. (checked by @leon)
                for (_symbol, py_book) in &mut self.books {
                    let mut book: PyRefMut<BookBuilder> = py_book.as_mut(py);
                    book.clear();
                }

                for entry in &record.data.data {
                    if !self.book_symbols.contains(&entry.symbol) {
                        continue;
                    }

                    let py_book: &mut Py<BookBuilder> = self
                        .books
                        .entry(entry.symbol.clone())
                        .or_insert_with(|| Py::new(py, BookBuilder::new()).unwrap());
                    let mut book: PyRefMut<BookBuilder> = py_book.as_mut(py);

                    match &*entry.side {
                        "Sell" => {
                            book.update_ask(entry.price, entry.size as f64);
                        }
                        "Buy" => {
                            book.update_bid(entry.price, entry.size as f64);
                        }
                        _ => panic!("Bitmex API changed"),
                    }
                    book.update_timestamp(record.timestamp);
                    self.id_to_price.insert(entry.id, entry.price);
                }
            }
            "insert" => {
                if !self.any_snapshot {
                    return Ok(());
                }
                for entry in record.data.data {
                    if !self.book_symbols.contains(&entry.symbol) {
                        continue;
                    }

                    let py_book: &mut Py<BookBuilder> =
                        &mut self.books.get_mut(&entry.symbol).unwrap();
                    let mut book: PyRefMut<BookBuilder> = py_book.as_mut(py);
                    self.id_to_price.insert(entry.id, entry.price);

                    match &*entry.side {
                        "Sell" => {
                            let price = self.id_to_price[&entry.id];
                            book.update_ask(price, entry.size as f64);
                            book.prune_crossing_with_ask(price);
                        }
                        "Buy" => {
                            let price = self.id_to_price[&entry.id];
                            book.update_bid(price, entry.size as f64);
                            book.prune_crossing_with_bid(price);
                        }
                        _ => panic!("Bitmex API changed"),
                    }
                }
            }
            "delete" => {
                if !self.any_snapshot {
                    return Ok(());
                }
                for entry in record.data.data {
                    if !self.book_symbols.contains(&entry.symbol) {
                        continue;
                    }

                    let py_book: &mut Py<BookBuilder> =
                        &mut self.books.get_mut(&entry.symbol).unwrap();
                    let mut book: PyRefMut<BookBuilder> = py_book.as_mut(py);

                    match &*entry.side {
                        "Sell" => {
                            book.delete_ask(self.id_to_price[&entry.id]);
                        }
                        "Buy" => {
                            book.delete_bid(self.id_to_price[&entry.id]);
                        }
                        _ => panic!("Bitmex API changed"),
                    }
                }
            }
            "update" => {
                if !self.any_snapshot {
                    return Ok(());
                }
                // {
                // "table":"orderBookL2_25",
                // "action":"update",
                // "data":[
                //     {"symbol":"XBTUSD","id":17999995000,"side":"Buy","size":5}
                // ]
                // }
                for entry in record.data.data {
                    if !self.book_symbols.contains(&entry.symbol) {
                        continue;
                    }

                    let py_book: &mut Py<BookBuilder> =
                        &mut self.books.get_mut(&entry.symbol).unwrap();
                    let mut book: PyRefMut<BookBuilder> = py_book.as_mut(py);

                    match &*entry.side {
                        "Sell" => {
                            let price = self.id_to_price[&entry.id];
                            book.update_ask(price, entry.size as f64);
                            book.prune_crossing_with_ask(price);
                        }
                        "Buy" => {
                            let price = self.id_to_price[&entry.id];
                            book.update_bid(price, entry.size as f64);
                            book.prune_crossing_with_bid(price);
                        }
                        _ => panic!("Bitmex API changed"),
                    }
                }
            }
            _ => panic!("Bitmex API changed"),
        }
        for symbol in touched_symbols {
            let py_book: &mut Py<BookBuilder> = &mut self.books.get_mut(&symbol).unwrap();
            {
                let mut book: PyRefMut<BookBuilder> = py_book.as_mut(py);
                book.update_timestamp(record.timestamp);
            }
            for callback in &self.book_callbacks {
                callback.call1(py, (symbol.clone(), py_book.to_object(py)))?;
            }
        }
        Ok(())
    }

    fn register_book_symbol(&mut self, symbol: &PyString) -> PyResult<()> {
        let symbol = symbol.to_string()?.into_owned();
        self.book_symbols.insert(symbol);
        Ok(())
    }
    fn register_book_callback(&mut self, callback: PyObject) {
        self.book_callbacks.push(callback);
    }
}
